"""
GMSH可视化模块
包含所有与GMSH相关的可视化函数
"""

import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from mpl_toolkits.mplot3d import Axes3D


def visualize_vtk_face_ids_matplotlib_2d(vtk_file_path, title="Face ID Visualization (2D)", fig=None, ax=None):
    """
    使用matplotlib二维显示VTK文件中的FaceID数据
    只显示XY平面，忽略Z坐标
    """
    try:
        import vtk
        from vtk.util import numpy_support
        
        # 读取VTK文件
        reader = vtk.vtkPolyDataReader()
        reader.SetFileName(vtk_file_path)
        reader.Update()
        
        polydata = reader.GetOutput()
        
        if polydata is None or polydata.GetNumberOfPoints() == 0:
            print(f"警告: VTK文件 {vtk_file_path} 不包含有效数据")
            return None, None
        
        # 创建二维图形
        if fig is None or ax is None:
            fig, ax = plt.subplots(figsize=(12, 8))
        
        # 获取点坐标
        points = polydata.GetPoints()
        points_array = numpy_support.vtk_to_numpy(points.GetData())
        
        # 获取FaceID数据
        face_id_array = polydata.GetCellData().GetArray("FaceID")
        
        if face_id_array is None:
            print(f"警告: VTK文件 {vtk_file_path} 中未找到FaceID数组")
            face_ids = np.zeros(polydata.GetNumberOfPolys())
        else:
            face_ids = numpy_support.vtk_to_numpy(face_id_array)
        
        # 获取多边形数据
        polys = polydata.GetPolys()
        polys_array = numpy_support.vtk_to_numpy(polys.GetData())
        
        # 为不同的FaceID分配颜色
        unique_face_ids = np.unique(face_ids)
        colors = plt.cm.tab20(np.linspace(0, 1, len(unique_face_ids)))
        
        face_id_to_color = {face_id: colors[i] for i, face_id in enumerate(unique_face_ids)}
        
        # 处理每个多边形
        idx = 0
        cell_index = 0
        
        while idx < len(polys_array):
            n_points = polys_array[idx]
            idx += 1
            
            if n_points >= 3:  # 至少需要3个点构成多边形
                polygon_indices = polys_array[idx:idx+n_points]
                idx += n_points
                
                # 获取多边形顶点坐标
                polygon_points = points_array[polygon_indices]
                
                # 获取FaceID
                face_id = face_ids[cell_index] if cell_index < len(face_ids) else 0
                cell_index += 1
                
                # 提取XY坐标（二维显示）
                x = polygon_points[:, 0]
                y = polygon_points[:, 1]
                
                # 绘制多边形
                polygon = patches.Polygon(np.column_stack([x, y]), 
                                        closed=True, 
                                        fill=True, 
                                        color=face_id_to_color.get(face_id, 'blue'), 
                                        alpha=0.7,
                                        edgecolor='black',
                                        linewidth=0.5)
                ax.add_patch(polygon)
        
        # 设置图形属性
        ax.set_title(title, fontsize=14)
        ax.set_xlabel('X坐标', fontsize=12)
        ax.set_ylabel('Y坐标', fontsize=12)
        ax.grid(True, alpha=0.3)
        ax.set_aspect('equal')
        
        # 添加图例
        legend_patches = []
        for face_id in unique_face_ids:
            patch = patches.Patch(color=face_id_to_color[face_id], label=f"Face ID {int(face_id)}")
            legend_patches.append(patch)
        
        if legend_patches:
            ax.legend(handles=legend_patches, loc='upper right', bbox_to_anchor=(1.15, 1))
        
        # 自动调整坐标轴范围
        ax.autoscale_view()
        
        print(f"{title} 已使用matplotlib二维显示")
        
        return fig, ax
        
    except Exception as e:
        print(f"处理VTK文件时出错: {e}")
        import traceback
        traceback.print_exc()
        return None, None


def visualize_vtk_face_ids_matplotlib(vtk_file_path, title="Face ID Visualization", fig=None, ax=None):
    """
    使用matplotlib可视化VTK/VTU文件中的FaceID数据
    支持VTK和VTU格式文件
    """
    # 检查文件扩展名
    file_ext = os.path.splitext(vtk_file_path)[1].lower()
    
    if file_ext == '.vtu':
        # 处理VTU文件
        try:
            import meshio
            
            # 读取VTU文件
            mesh = meshio.read(vtk_file_path)
            
            # 创建图形
            if fig is None or ax is None:
                fig = plt.figure(figsize=(12, 8))
                ax = fig.add_subplot(111, projection='3d')
            
            # 获取点坐标
            points = mesh.points
            
            # 获取FaceID数据
            face_ids = None
            if 'FaceID' in mesh.cell_data:
                face_ids = mesh.cell_data['FaceID'][0]
            
            # 处理四边形单元
            for cell_block in mesh.cells:
                if cell_block.type == 'quad':
                    for cell in cell_block.data:
                        # 获取四边形的四个顶点
                        quad_points = points[cell]
                        
                        # 获取对应的FaceID
                        face_id = 0
                        if face_ids is not None:
                            # 这里需要根据单元索引获取FaceID
                            # 简化处理：使用第一个FaceID
                            face_id = face_ids[0] if len(face_ids) > 0 else 0
                        
                        # 绘制四边形（分解为两个三角形）
                        x = quad_points[:, 0]
                        y = quad_points[:, 1]
                        z = quad_points[:, 2]
                        
                        # 绘制第一个三角形
                        ax.plot_trisurf([x[0], x[1], x[2]], [y[0], y[1], y[2]], [z[0], z[1], z[2]], 
                                       color=plt.cm.tab20(face_id % 20), alpha=0.7)
                        # 绘制第二个三角形
                        ax.plot_trisurf([x[0], x[2], x[3]], [y[0], y[2], y[3]], [z[0], z[2], z[3]], 
                                       color=plt.cm.tab20(face_id % 20), alpha=0.7)
            
            # 设置图形属性
            ax.set_title(title)
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_zlabel('Z')
            
            # 添加颜色条
            if face_ids is not None:
                unique_face_ids = np.unique(face_ids)
                if len(unique_face_ids) > 1:
                    sm = plt.cm.ScalarMappable(cmap=plt.cm.tab20, 
                                            norm=plt.Normalize(vmin=min(unique_face_ids), 
                                                            vmax=max(unique_face_ids)))
                    sm.set_array([])
                    cbar = plt.colorbar(sm, ax=ax, shrink=0.8, aspect=20)
                    cbar.set_label('Face ID')
            
            print(f"{title} 已使用matplotlib显示")
            
        except ImportError:
            print("错误: meshio库未安装，无法读取VTU文件")
            return None, None


def visualize_gmsh_physical_groups_2d(msh_file_path, title="GMSH物理组可视化", fig=None, ax=None):
    """
    使用matplotlib二维显示GMSH msh文件的物理组颜色映射
    直接读取msh文件，无需转换为VTK格式
    """
    import gmsh
    
    # 初始化Gmsh
    gmsh.initialize()
    gmsh.option.setNumber("General.Terminal", 0)  # 关闭终端输出
    
    try:
        # 加载msh文件
        gmsh.open(msh_file_path)
        print(f"已加载GMSH网格文件: {msh_file_path}")
        
        # 获取所有物理组
        physical_groups = gmsh.model.getPhysicalGroups()
        print(f"找到 {len(physical_groups)} 个物理组")
        
        # 创建图形（二维显示）
        if fig is None or ax is None:
            fig, ax = plt.subplots(figsize=(12, 8))
        
        # 为每个物理组分配颜色
        colors = plt.cm.tab20(np.linspace(0, 1, len(physical_groups)))
        
        # 处理每个物理组
        for i, (dim, tag) in enumerate(physical_groups):
            # 获取物理组名称
            name = gmsh.model.getPhysicalName(dim, tag)
            print(f"物理组 {i}: 维度={dim}, 标签={tag}, 名称='{name}'")
            
            # 获取物理组包含的实体
            entities = gmsh.model.getEntitiesForPhysicalGroup(dim, tag)
            
            # 处理每个实体
            for entity in entities:
                # 获取实体上的单元
                element_types, element_tags, element_node_tags = gmsh.model.mesh.getElements(dim, entity)
                
                # 处理每种单元类型
                for elem_type, elem_tags, elem_node_tags in zip(element_types, element_tags, element_node_tags):
                    # 只处理面单元（三角形和四边形）
                    if elem_type in [2, 3]:  # 2=三角形, 3=四边形
                        print(f"处理单元类型 {elem_type}, 单元数量: {len(elem_tags)}")
                        
                        # 根据单元类型确定每个单元的节点数
                        if elem_type == 2:  # 三角形，3个节点
                            nodes_per_element = 3
                        elif elem_type == 3:  # 四边形，4个节点
                            nodes_per_element = 4
                        
                        # 处理每个单元
                        for j in range(len(elem_tags)):
                            # 计算节点标签在elem_node_tags数组中的起始位置
                            start_idx = j * nodes_per_element
                            end_idx = start_idx + nodes_per_element
                            
                            if end_idx <= len(elem_node_tags):
                                # 获取该单元的节点标签
                                node_tags = elem_node_tags[start_idx:end_idx]
                                
                                # 获取节点坐标
                                node_coords = []
                                for node_tag in node_tags:
                                    coord_result = gmsh.model.mesh.getNode(node_tag)
                                    if coord_result is not None:
                                        # coord_result是一个元组，第一个元素是坐标数组
                                        coord_array = coord_result[0]
                                        if len(coord_array) == 3:
                                            node_coords.append(coord_array)
                                        else:
                                            # 如果坐标不是三维，填充为三维
                                            if len(coord_array) == 2:
                                                node_coords.append([coord_array[0], coord_array[1], 0.0])
                                            else:
                                                node_coords.append([0.0, 0.0, 0.0])
                                
                                if len(node_coords) == 0:
                                    continue
                                    
                                # 转换为numpy数组，确保所有坐标都是三维的
                                node_coords = np.array(node_coords)
                                
                                # 提取XY坐标（二维显示）
                                x = node_coords[:, 0]
                                y = node_coords[:, 1]
                                
                                # 绘制单元
                                if elem_type == 2:  # 三角形
                                    triangle = patches.Polygon(np.column_stack([x, y]), 
                                                             closed=True, 
                                                             fill=True, 
                                                             color=colors[i], 
                                                             alpha=0.7,
                                                             edgecolor='black',
                                                             linewidth=0.5)
                                    ax.add_patch(triangle)
                                
                                elif elem_type == 3:  # 四边形
                                    quad = patches.Polygon(np.column_stack([x, y]), 
                                                          closed=True, 
                                                          fill=True, 
                                                          color=colors[i], 
                                                          alpha=0.7,
                                                          edgecolor='black',
                                                          linewidth=0.5)
                                    ax.add_patch(quad)
        
        # 设置图形属性
        ax.set_title(title, fontsize=14)
        ax.set_xlabel('X坐标')
        ax.set_ylabel('Y坐标')
        ax.grid(True, alpha=0.3)
        ax.set_aspect('equal')
        
        # 创建图例
        legend_patches = []
        for i, (dim, tag) in enumerate(physical_groups):
            name = gmsh.model.getPhysicalName(dim, tag)
            patch = patches.Patch(color=colors[i], label=f'{name} (dim={dim}, tag={tag})')
            legend_patches.append(patch)
        
        ax.legend(handles=legend_patches, loc='upper right', bbox_to_anchor=(1.15, 1))
        
        # 自动调整坐标轴范围
        ax.autoscale_view()
        
        print(f"{title} 已使用matplotlib二维显示")
        
        return fig, ax
        
    except Exception as e:
        print(f"GMSH二维可视化失败: {e}")
        import traceback
        traceback.print_exc()
        return None, None
    
    finally:
        # 清理Gmsh
        gmsh.finalize()


def compare_gmsh_vtk_visualization(msh_file_path, vtk_file_path=None):
    """
    同时显示GMSH物理组可视化和VTK FaceID可视化进行对比
    """
    import os
    
    # 创建子图布局
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))
    
    print("=" * 60)
    print("开始对比可视化")
    print("=" * 60)
    
    # 显示GMSH物理组可视化
    if os.path.exists(msh_file_path):
        print(f"✓ 加载GMSH文件: {msh_file_path}")
        fig_gmsh, ax_gmsh = visualize_gmsh_physical_groups_2d(msh_file_path, "GMSH物理组可视化", fig, ax1)
        if fig_gmsh is not None:
            print("✓ GMSH可视化完成")
        else:
            print("✗ GMSH可视化失败")
            ax1.text(0.5, 0.5, 'GMSH可视化失败', ha='center', va='center', transform=ax1.transAxes)
    else:
        print(f"✗ GMSH文件不存在: {msh_file_path}")
        ax1.text(0.5, 0.5, 'GMSH文件不存在', ha='center', va='center', transform=ax1.transAxes)
    
    # 显示VTK FaceID可视化
    if vtk_file_path and os.path.exists(vtk_file_path):
        print(f"✓ 加载VTK文件: {vtk_file_path}")
        fig_vtk, ax_vtk = visualize_vtk_face_ids_matplotlib_2d(vtk_file_path, "VTK FaceID可视化", fig, ax2)
        if fig_vtk is not None:
            print("✓ VTK可视化完成")
        else:
            print("✗ VTK可视化失败")
            ax2.text(0.5, 0.5, 'VTK可视化失败', ha='center', va='center', transform=ax2.transAxes)
    else:
        print("✗ VTK文件未提供或不存在")
        ax2.text(0.5, 0.5, 'VTK文件未提供', ha='center', va='center', transform=ax2.transAxes)
    
    # 调整布局
    plt.tight_layout()
    
    print("=" * 60)
    print("对比可视化完成，正在显示图形...")
    print("=" * 60)
    
    # 显示图形
    plt.show()
    
    return fig, (ax1, ax2)